package com.ktjiaoyu.travelaiagent.controller;

import cn.hutool.core.util.ArrayUtil;
import com.ktjiaoyu.travelaiagent.agent.KtTravelManus;
import com.ktjiaoyu.travelaiagent.app.TravelApp;
import com.ktjiaoyu.travelaiagent.mcp.McpToolCallbackProxy;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.Arrays;
import java.util.UUID;

/**
 * 与AI聊天的控制器
 */
@RestController
@RequestMapping("/ai")
public class ChatWithAIController {
    // 注入TravelApp实例
    @Resource
    private TravelApp travelApp;
    /**
     * 前端可以通过此方法获得一个聊天会话ID
     */
    @GetMapping("/chat/new")
    public String newChat() {
        // 返回一个新的会话
        return UUID.randomUUID().toString();
    }
    /**
     * 与AI聊天（同步）
     */
    @GetMapping("/chat/travel/sync")
    public String chatWithTravelAppSync(String chatId, String message) {
        return travelApp.chat(chatId, message);
    }

    /**
     * SSE实现方式一：返回Flux对象
     * 必须设置produces为MediaType.TEXT_EVENT_STREAM_VALUE
     */
    @GetMapping(value = "/chat/travel/sse/flux", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chatWithTravelAppSSE_Flux(String chatId, String message) {
        return travelApp.streamChat(chatId, message);
    }
    /**
     * SSE实现方式二：使用ServerSentEvent泛型
     */
    @GetMapping(value = "/chat/travel/sse/generic")
    public Flux<ServerSentEvent<String>> chatWithTravelAppSSE_ServerSentEvent(String chatId, String message) {
        return travelApp.streamChat(chatId, message)
                // 将字符串片段包装成 ServerSentEvent对象
                .map(chunk -> ServerSentEvent.<String>builder()
                        .data(chunk)
                        .build());
    }

    /**
     * SSE实现方式三：使用SseEmitter
     */
    @GetMapping("/chat/travel/sse/emitter")
    public SseEmitter chatWithTravelAppSSE_SseEmitter(String chatId, String message) {
        // 创建一个超时时间较长的 SseEmitter (根据实际情况，若有工具调用要更长一些)
        SseEmitter emitter = new SseEmitter(180000L); // 3分钟应该要输出完成，否则将超时
        // 获取 Flux 数据流
        travelApp.streamChat(chatId, message)
                // 订阅消息
                .subscribe(
                        // 处理每条消息
                        chunk -> {
                            try {
                                // 收到消息就发送
                                emitter.send(chunk);
                            } catch (IOException e) {
                                emitter.completeWithError(e);
                            }
                        },
                        // 处理错误
                        emitter::completeWithError,
                        // 处理完成
                        emitter::complete
                );
        // 返回emitter
        return emitter;
    }

    @Resource
    private ToolCallback[] aiTools;
    @Resource
    private ChatModel dashscopeChatModel;
    @Resource
    private ToolCallbackProvider toolCallbackProvider;
    /**
     * 流式调用 Manus 智能体
     * @param message
     * @return
     */
    @GetMapping("/chat/manus/sse")
    public SseEmitter chatWithManus(String message) {
        // 获得合并后的工具数组
        ToolCallback[] allTools = mergeAiTools();
        // 暂时不考虑恢复会话，每次请求都是新的智能体（注意这里改成传递全部工具）
        KtTravelManus ktTravelManus = new KtTravelManus(allTools, dashscopeChatModel);
        // 设置最大步数，这个最好是通过配置设置
        ktTravelManus.setMaxSteps(20);
        return ktTravelManus.streamRun(message);
    }
    /*
     * 合并本地工具和MCP工具，组成一个新的数组
     */
    private ToolCallback[] mergeAiTools() {
        // 注意要进行代理，避免MCP工具出现ToolContext问题
        FunctionCallback[] mcpFunctionCallbacks = McpToolCallbackProxy
                .proxyAll(toolCallbackProvider.getToolCallbacks());
        // 转化为ToolCallback数组
        ToolCallback[] mcpToolCallbacks =
                Arrays.stream(mcpFunctionCallbacks)
                        .map(tool -> (ToolCallback) tool)
                        .toArray(ToolCallback[]::new);
        // 合并本地工具和MCP工具
        ToolCallback[] allTools = ArrayUtil.addAll(aiTools, mcpToolCallbacks);
        return allTools;
    }

}
